#!/usr/bin/env python3
# 青龙面板内运行的多任务日志迁移脚本（修复版）
# cron: 2 0 * * *
# new Env("青龙日志迁移工具")
import os
import paramiko
import sys
import traceback
import time
from datetime import datetime

def log(message, level="INFO"):
    """记录带时间戳的日志"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{level}] {message}")
    sys.stdout.flush()

def migrate_logs(ssh_host, ssh_port, ssh_username, ssh_password, source_container, target_container):
    """执行单个迁移任务"""
    # 固定日志路径
    log_path = "/ql/data/log/Message/"
    
    try:
        log(f"开始迁移任务: {source_container} -> {target_container}")
        
        # 创建SSH客户端
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(ssh_host, port=ssh_port, username=ssh_username, password=ssh_password)
        
        # 1. 验证源容器
        stdin, stdout, stderr = ssh.exec_command(
            f"docker exec {source_container} ls {log_path}", 
            get_pty=True
        )
        if stdout.channel.recv_exit_status() != 0:
            log(f"源容器或路径无效: {stderr.read().decode()}", "ERROR")
            return False
        
        # 2. 准备目标目录（确保存在）
        ssh.exec_command(
            f"docker exec {target_container} mkdir -p {log_path}", 
            get_pty=True
        )
        
        # 3. 修复：仅迁移文件（不保留源路径结构）
        log("正在迁移文件...")
        # 关键修改：使用通配符 * 直接复制文件
        cmd = f"docker exec {source_container} sh -c 'cd {log_path} && tar -cf - *' | " \
              f"docker exec -i {target_container} tar -xf - -C {log_path}"
        
        stdin, stdout, stderr = ssh.exec_command(cmd, get_pty=True)
        exit_status = stdout.channel.recv_exit_status()
        if exit_status != 0:
            error = stderr.read().decode().strip()
            log(f"迁移失败 (状态码 {exit_status}): {error}", "ERROR")
            return False
        
        # 4. 修复权限
        ssh.exec_command(
            f"docker exec {target_container} chown -R 1000:1000 {log_path}", 
            get_pty=True
        )
        
        # 5. 验证结果
        stdin, stdout, stderr = ssh.exec_command(
            f"docker exec {target_container} ls -lh {log_path} | head -5", 
            get_pty=True
        )
        result = stdout.read().decode().strip()
        log(f"迁移结果:\n{result}")
        
        log(f"✓ {source_container} -> {target_container} 迁移成功")
        return True
        
    except Exception as e:
        log(f"❌ 迁移失败: {str(e)}", "ERROR")
        log(traceback.format_exc(), "DEBUG")
        return False
    finally:
        try:
            ssh.close()
        except:
            pass

def main():
    # 安装依赖（只需运行一次）
    try:
        import paramiko
    except ImportError:
        log("正在安装paramiko...")
        os.system("pip install paramiko -i https://pypi.tuna.tsinghua.edu.cn/simple")
        import paramiko
    
    # 读取环境变量
    migrate_config = os.getenv("Migrate_RZQY", "")
    if not migrate_config:
        log("❌ 未找到迁移配置环境变量 Migrate_RZQY")
        return
    
    log("="*70)
    log("青龙日志迁移工具")
    log(f"开始时间: {datetime.now()}")
    log(f"找到 {len(migrate_config.splitlines())} 个迁移任务")
    log("="*70)
    
    # 解析并执行迁移任务
    success_count = 0
    for i, line in enumerate(migrate_config.splitlines()):
        if not line.strip() or line.startswith("#"):
            continue
            
        log(f"\n任务 #{i+1}: {line}")
        parts = line.split("#")
        if len(parts) != 6:
            log("❌ 配置格式错误，需要 SSH_HOST#SSH_PORT#SSH_USERNAME#SSH_PASSWORD#SOURCE_CONTAINER#TARGET_CONTAINER")
            continue
            
        ssh_host, ssh_port, ssh_username, ssh_password, source_container, target_container = parts
        
        # 执行迁移
        if migrate_logs(ssh_host, int(ssh_port), ssh_username, ssh_password, source_container, target_container):
            success_count += 1
    
    # 结果总结
    log("\n" + "="*70)
    log(f"迁移完成! 成功: {success_count}/{len(migrate_config.splitlines())}")
    log("="*70)
    
    if success_count == 0:
        sys.exit(1)

if __name__ == "__main__":
    main()